import numpy as np
from SE3 import Pose, quaternionSlerp

class MotionPlanner:
    def __init__(self, max_waypoint_num):
        assert max_waypoint_num > 0
        self.max_waypoint_num = max_waypoint_num
        self.joints_step = np.deg2rad([5, 4, 4, 3, 3, 3])
        self.linear_step = 5
        pass
    
    #privarte
    def __getMid(self, start_value, end_value, t):
        mid_value = (1-t)*start_value + t*end_value
        return mid_value
    
    #public method
    def joint_space_planner(self, start_joints, end_joints):
        ''' motion plann in joint_space(j1, j2, j3, j4, j5, j6)
        '''
        mid_joints = []
        cnt1 = int(np.floor(np.fabs(start_joints[0] - end_joints[0]) / self.joints_step[0])) + 1
        cnt2 = int(np.floor(np.fabs(start_joints[1] - end_joints[1]) / self.joints_step[1])) + 1
        cnt3 = int(np.floor(np.fabs(start_joints[2] - end_joints[2]) / self.joints_step[2])) + 1
        cnt4 = int(np.floor(np.fabs(start_joints[3] - end_joints[3]) / self.joints_step[3])) + 1
        cnt5 = int(np.floor(np.fabs(start_joints[4] - end_joints[4]) / self.joints_step[4])) + 1
        cnt6 = int(np.floor(np.fabs(start_joints[5] - end_joints[5]) / self.joints_step[5])) + 1
        cnt_max = max([cnt1, cnt2, cnt3, cnt4, cnt5, cnt6])
        cnt_max = self.max_waypoint_num if cnt_max > self.max_waypoint_num else cnt_max
        for i in range(cnt_max+1):
            t = 1.0*i/cnt_max
            joints = []
            joints.append(self.__getMid(start_joints[0], end_joints[0], t))
            joints.append(self.__getMid(start_joints[1], end_joints[1], t))
            joints.append(self.__getMid(start_joints[2], end_joints[2], t))
            joints.append(self.__getMid(start_joints[3], end_joints[3], t))
            joints.append(self.__getMid(start_joints[4], end_joints[4], t))
            joints.append(self.__getMid(start_joints[5], end_joints[5], t))
            mid_joints.append(joints)
        return mid_joints
    
    def cartesian_space_linear_planner(self, start_pose: Pose, end_pose: Pose):
        '''linear position and quaterion interpolation algorithm
        '''
        # interpolation accuracy: 10mm
        way_points = []
        start_quat = [start_pose.qx, start_pose.qy, start_pose.qz, start_pose.qw]
        end_quat = [end_pose.qx, end_pose.qy, end_pose.qz, end_pose.qw]
        dist = np.sqrt(pow((end_pose.x - start_pose.x), 2) + 
                        pow((end_pose.y - start_pose.y), 2) + 
                        pow((end_pose.z - start_pose.z), 2))
        cnt = int(dist / self.linear_step) + 1
        cnt = self.max_waypoint_num if cnt > self.max_waypoint_num else cnt
        for i in range(cnt+1):
            t = float((1.0 * i) / cnt)
            temp_pose = Pose(0, 0, 0, 0, 0, 0, 1)
            temp_pose.x = self.__getMid(start_pose.x, end_pose.x, t)
            temp_pose.y = self.__getMid(start_pose.y, end_pose.y, t)
            temp_pose.z = self.__getMid(start_pose.z, end_pose.z, t)
            quat = quaternionSlerp(start_quat, end_quat, t)
            temp_pose.qx = quat[0]
            temp_pose.qy = quat[1]
            temp_pose.qz = quat[2]
            temp_pose.qw = quat[3]
            way_points.append(temp_pose)
        return way_points
    
    def cartesian_space_circle_planner(self, start_pose: Pose, mid_pose: Pose, end_pose: Pose):
        '''circle position and quaterion interpolation algorithm
        '''
        pass
    
if __name__ == '__main__':
    mp = MotionPlanner(10)
    pose1 = Pose(0, 0, 0, 0, 0, 0, 1)
    pose2 = Pose(20, 30, 40, 0, 0, 0.7071, 0.7071)
    way_points = mp.cartesian_space_linear_planner(pose1, pose2)
    for i in range(len(way_points)):
        temp_pose = way_points[i]
        print(f"=======================================")
        print(f"pose{i}----------pose.x = {temp_pose.x}")
        print(f"pose{i}----------pose.y = {temp_pose.y}")
        print(f"pose{i}----------pose.z = {temp_pose.z}")
        print(f"pose{i}----------pose.qx = {temp_pose.qx}")
        print(f"pose{i}----------pose.qy = {temp_pose.qy}")
        print(f"pose{i}----------pose.qz = {temp_pose.qz}")
        print(f"pose{i}----------pose.qw = {temp_pose.qw}")
    pass


